#!/usr/bin/env python3
# alloc_pages_monitor.py

from bcc import BPF
import ctypes as ct
import time
import argparse
from collections import defaultdict, deque
import struct

# 定义与BPF程序中相同的数据结构
class AllocEvent(ct.Structure):
    _fields_ = [
        ("pid", ct.c_uint32),
        ("tid", ct.c_uint32),
        ("timestamp", ct.c_uint64),
        ("gfp_flags", ct.c_uint64),
        ("order", ct.c_uint32),
        ("page_addr", ct.c_uint64),
        ("nr_pages", ct.c_uint32),
        ("comm", ct.c_char * 16)  # TASK_COMM_LEN
    ]

class PageAllocMonitor:
    def __init__(self, target_pid=None, duration=30):
        self.bpf = None
        self.target_pid = target_pid
        self.duration = duration
        
        # 统计数据结构
        self.stats = {
            'total_allocations': 0,
            'total_pages': 0,
            'by_order': defaultdict(int),
            'by_process': defaultdict(lambda: {'count': 0, 'pages': 0}),
            'gfp_flags_stats': defaultdict(int),
            'timeline': deque(maxlen=10000)
        }
        
        self.load_bpf_program()

    def load_bpf_program(self):
        """加载BPF程序"""
        # self.bpf = BPF(text=bpf_text)
        self.bpf = BPF(src_file="memory.c")

        # 挂载探针
        self.bpf.attach_kprobe(
            event="alloc_pages_noprof", 
            fn_name="trace_alloc_pages_entry"
        )
        self.bpf.attach_kretprobe(
            event="alloc_pages_noprof", 
            fn_name="trace_alloc_pages_return"
        )

    def decode_gfp_flags(self, gfp_flags):
        """解码GFP标志"""
        flags = []
        gfp_map = {
            (1 << 0): "__GFP_DMA",
            (1 << 1): "__GFP_HIGHMEM", 
            (1 << 2): "__GFP_DMA32",
            (1 << 3): "__GFP_MOVABLE",
            (1 << 4): "__GFP_RECLAIMABLE",
            (1 << 5): "__GFP_HIGH",
            (1 << 6): "__GFP_IO",
            (1 << 7): "__GFP_FS",
            (1 << 8): "__GFP_ZERO",
            (1 << 9): "__GFP_NOWARN",
        }
        
        for flag, name in gfp_map.items():
            if gfp_flags & flag:
                flags.append(name)
        
        return "|".join(flags) if flags else "GFP_KERNEL"

    def process_event(self, cpu, data, size):
        """处理分配事件"""
        event = self.bpf["alloc_events"].event(data)
        
        # 更新统计信息
        self.stats['total_allocations'] += 1
        self.stats['total_pages'] += event.nr_pages
        self.stats['by_order'][event.order] += 1
        
        process_key = f"{event.comm.decode()}[{event.pid}]"
        self.stats['by_process'][process_key]['count'] += 1
        self.stats['by_process'][process_key]['pages'] += event.nr_pages
        
        gfp_str = self.decode_gfp_flags(event.gfp_flags)
        self.stats['gfp_flags_stats'][gfp_str] += 1
        
        # 添加到时间线
        timeline_event = {
            'timestamp': event.timestamp,
            'process': process_key,
            'order': event.order,
            'pages': event.nr_pages,
            'gfp_flags': gfp_str,
            'page_addr': event.page_addr
        }
        self.stats['timeline'].append(timeline_event)
        
        # 实时显示（可选）
        if self.stats['total_allocations'] % 100 == 0:
            self.display_realtime_stats()

    def display_realtime_stats(self):
        """显示实时统计"""
        print(f"\r分配次数: {self.stats['total_allocations']} | "
              f"总页数: {self.stats['total_pages']} | "
              f"活跃进程: {len(self.stats['by_process'])}", end="")

    def generate_report(self):
        """生成详细报告"""
        print("\n" + "="*80)
        print("alloc_pages_noprof 跟踪报告")
        print("="*80)
        
        print(f"\n总体统计:")
        print(f"  总分配次数: {self.stats['total_allocations']}")
        print(f"  总分配页数: {self.stats['total_pages']}")
        print(f"  平均每次分配页数: {self.stats['total_pages'] / max(1, self.stats['total_allocations']):.2f}")
        
        print(f"\n按分配阶数统计:")
        for order, count in sorted(self.stats['by_order'].items()):
            pages = (1 << order) * count
            print(f"  order {order} ({1 << order}页): {count}次, {pages}页")
        
        print(f"\n按进程统计:")
        for process, data in sorted(self.stats['by_process'].items(), 
                                  key=lambda x: x[1]['pages'], reverse=True)[:10]:
            print(f"  {process}: {data['count']}次分配, {data['pages']}页")
        
        print(f"\nGFP标志使用统计:")
        for gfp, count in sorted(self.stats['gfp_flags_stats'].items(), 
                               key=lambda x: x[1], reverse=True)[:5]:
            print(f"  {gfp}: {count}次")

    def run(self):
        """运行监控"""
        print("开始跟踪 alloc_pages_noprof 调用...")
        print("按 Ctrl+C 停止跟踪")
        
        # 注册事件处理
        self.bpf["alloc_events"].open_perf_buffer(self.process_event)
        
        start_time = time.time()
        try:
            while time.time() - start_time < self.duration:
                self.bpf.perf_buffer_poll(timeout=1000)
        except KeyboardInterrupt:
            print("\n跟踪被用户中断")
        
        self.generate_report()

def main():
    parser = argparse.ArgumentParser(description='跟踪 alloc_pages_noprof 函数')
    parser.add_argument('-p', '--pid', type=int, help='跟踪特定进程')
    parser.add_argument('-d', '--duration', type=int, default=30, 
                       help='跟踪时长（秒）')
    parser.add_argument('-v', '--verbose', action='store_true',
                       help='显示详细输出')
    
    args = parser.parse_args()
    
    monitor = PageAllocMonitor(target_pid=args.pid, duration=args.duration)
    monitor.run()

if __name__ == "__main__":
    main()